package transport

import (
	"bytes"
	"fmt"
	"net"
	"rtps/common"
	"rtps/messages"
)

// UDP传输层
type UDPTransport struct {
	unicastConn   *net.UDPConn
	multicastConn *net.UDPConn
	localAddr     *net.UDPAddr
	multicastAddr *net.UDPAddr
}

// 创建新的UDP传输实例
func NewUDPTransport(port int) (*UDPTransport, error) {
	transport := &UDPTransport{}

	// 创建单播连接
	localAddr, err := net.ResolveUDPAddr("udp", fmt.Sprintf(":%d", port))
	if err != nil {
		return nil, fmt.Errorf("failed to resolve local address: %v", err)
	}

	unicastConn, err := net.ListenUDP("udp", localAddr)
	if err != nil {
		return nil, fmt.Errorf("failed to create unicast connection: %v", err)
	}

	// 创建多播连接
	multicastAddr, err := net.ResolveUDPAddr("udp", fmt.Sprintf("%s:%d",
		common.MulticastAddressSPDP, common.DefaultPortSPDP))
	if err != nil {
		unicastConn.Close()
		return nil, fmt.Errorf("failed to resolve multicast address: %v", err)
	}

	multicastConn, err := net.ListenMulticastUDP("udp", nil, multicastAddr)
	if err != nil {
		unicastConn.Close()
		return nil, fmt.Errorf("failed to create multicast connection: %v", err)
	}

	transport.unicastConn = unicastConn
	transport.multicastConn = multicastConn
	transport.localAddr = localAddr
	transport.multicastAddr = multicastAddr

	return transport, nil
}

// 发送消息到指定地址
func (t *UDPTransport) SendMessage(msg *messages.Message, addr *net.UDPAddr) error {
	// 序列化消息
	data, err := serializeMessage(msg)
	if err != nil {
		return fmt.Errorf("failed to serialize message: %v", err)
	}

	// 发送数据
	_, err = t.unicastConn.WriteToUDP(data, addr)
	if err != nil {
		return fmt.Errorf("failed to send message: %v", err)
	}

	return nil
}

// 发送多播消息
func (t *UDPTransport) SendMulticastMessage(msg *messages.Message) error {
	// 序列化消息
	data, err := serializeMessage(msg)
	if err != nil {
		return fmt.Errorf("failed to serialize message: %v", err)
	}

	// 发送多播数据
	_, err = t.multicastConn.WriteToUDP(data, t.multicastAddr)
	if err != nil {
		return fmt.Errorf("failed to send multicast message: %v", err)
	}

	return nil
}

// 接收消息
func (t *UDPTransport) ReceiveMessage() (*messages.Message, *net.UDPAddr, error) {
	buffer := make([]byte, 65507) // UDP最大数据包大小

	// 从单播连接接收
	n, addr, err := t.unicastConn.ReadFromUDP(buffer)
	if err != nil {
		return nil, nil, fmt.Errorf("failed to receive message: %v", err)
	}

	// 反序列化消息
	msg, err := deserializeMessage(buffer[:n])
	if err != nil {
		return nil, addr, fmt.Errorf("failed to deserialize message: %v", err)
	}

	return msg, addr, nil
}

// 接收多播消息
func (t *UDPTransport) ReceiveMulticastMessage() (*messages.Message, *net.UDPAddr, error) {
	buffer := make([]byte, 65507)

	// 从多播连接接收
	n, addr, err := t.multicastConn.ReadFromUDP(buffer)
	if err != nil {
		return nil, nil, fmt.Errorf("failed to receive multicast message: %v", err)
	}

	// 反序列化消息
	msg, err := deserializeMessage(buffer[:n])
	if err != nil {
		return nil, addr, fmt.Errorf("failed to deserialize multicast message: %v", err)
	}

	return msg, addr, nil
}

// 关闭传输层
func (t *UDPTransport) Close() error {
	var errs []error

	if t.unicastConn != nil {
		if err := t.unicastConn.Close(); err != nil {
			errs = append(errs, err)
		}
	}

	if t.multicastConn != nil {
		if err := t.multicastConn.Close(); err != nil {
			errs = append(errs, err)
		}
	}

	if len(errs) > 0 {
		return fmt.Errorf("errors closing transport: %v", errs)
	}

	return nil
}

// 序列化消息
func serializeMessage(msg *messages.Message) ([]byte, error) {
	// 使用bytes.Buffer进行序列化
	var buf bytes.Buffer
	if err := msg.Write(&buf); err != nil {
		return nil, err
	}

	return buf.Bytes(), nil
}

// 反序列化消息
func deserializeMessage(data []byte) (*messages.Message, error) {
	msg := &messages.Message{}
	buf := bytes.NewBuffer(data)

	if err := msg.Read(buf); err != nil {
		return nil, err
	}

	return msg, nil
}
